import math
import os

from constants import *
from data_process.path_algorithm.PathAlgorithm import PathAlgorithm
from utils.file_util import *

class MultiDirDiffAlgorithm(PathAlgorithm):
    def __init__(self, params):
        self.params = params
        self.rootPath = params.rootPath
        self.volun = params.volun
        self.innerDirName = params.innerDirName
        self.vocabPath = params.vocabPath
        self.mapPath = params.mapPath
        self.trainDirList = params.trainDirList
        self.testDirList = params.testDirList
        self.isGeneratePath = params.isGeneratePath

        self.trainPath = TRAIN_PATH
        self.testPath = TEST_PATH
        self.fileFormat = '*.wav'

    def generateTimeMarkFile(self):
        trainMarkPath = os.path.join(self.trainDir, TIME_MARK_PATH)
        testMarkPath = os.path.join(self.testDir, TIME_MARK_PATH)
        with open(trainMarkPath, 'rb') as f:
            trainMarkData = pickle.load(f)
        with open(testMarkPath, 'rb') as f:
            testMarkData = pickle.load(f)
        timeMarkMap = {}
        for key in trainMarkData.keys():
            savePath = os.path.join(self.trainDir, key)
            timeMarkMap[savePath] = trainMarkData[key]
        for key in testMarkData.keys():
            savePath = os.path.join(self.testDir, key)
            timeMarkMap[savePath] = testMarkData[key]

        maxLen = self.getCropMaxLen(timeMarkMap)
        timeMarkMap = self.changeMarkData(timeMarkMap, maxLen)

        with open('temp_path.txt', 'w') as f:
            for key in timeMarkMap.keys():
                f.write(f'{key} {timeMarkMap[key][0]} {timeMarkMap[key][1]}\n')

        with open(TOTAL_MARK_PATH, 'wb') as f:
            pickle.dump(timeMarkMap, f)

    def generatePath(self):
        self.createTrainPathFile()
        self.createTestPathFile()

        mapTabel = CreateMapTable(self.mapPath)
        createContentFromFilePath(self.trainPath, self.testPath, mapTabel)

    def addPrefixOfPathFile(self, filePath):
        savePaths = []
        prefix = self.trainDir
        paths = readDataFromFile(filePath)
        for curPath in paths:
            res = os.path.join(prefix, curPath)
            savePaths.append(res)
        with open(filePath, 'w') as f:
            for curPath in savePaths:
                f.write(f'{curPath}\n')


    def createTrainPathFile(self, ratioNum=-1):
        allPath = self.getAllPathFromAllDirList(self.trainDirList)
        leftPaths = self.selectPathWithRatio(allPath)
        leftPaths.sort()
        self.savePathToFile(self.trainPath, leftPaths)

    def createTestPathFile(self, ratioNum=-1):
        allPath = self.getAllPathFromAllDirList(self.testDirList)
        trainPaths = readDataFromFile(self.trainPath)
        vocab = self.getVocabFromPaths(trainPaths)
        leftPaths = self.filterPaths(allPath, vocab)
        leftPaths = self.selectPathWithRatio(leftPaths, 6)
        leftPaths.sort()
        self.savePathToFile(self.testPath, leftPaths)

    def selectPathWithRatio(self, paths, ratioNum=-1):
        vocab = self.getVocabFromPaths(paths)
        pathDict = {}
        for curPath in paths:
            prefix = self.getPrefixFromPath(curPath)
            if prefix not in pathDict.keys():
                pathDict[prefix] = []
            pathDict[prefix].append(curPath)
        resPath = []
        for key in pathDict.keys():
            curList = pathDict[key]
            if ratioNum == -1:
                selectNum = len(curList)
            else:
                selectNum = min(len(curList), ratioNum)
            resPath += curList[0:selectNum]
        return resPath

    def shuffleSplitPath(self, ratio):
        allTrains = readDataFromFile(self.trainPath)
        allTests = readDataFromFile(self.testPath)
        total = allTests + allTrains
        pathDict = {}
        for curPath in total:
            word = self.getWordFromPath(curPath)
            if word not in pathDict.keys():
                pathDict[word] = []
            pathDict[word].append(curPath)
        finalTrain = []
        finalTest = []
        for key in pathDict.keys():
            curList = pathDict[key]
            selectNum = math.ceil(len(curList) * ratio)
            sList = random.sample(range(0, len(curList)-1), selectNum)
            for i in range(len(curList)):
                if i in sList:
                    finalTest.append(curList[i])
                else:
                    finalTrain.append(curList[i])
        with open(self.trainPath, 'w') as f:
            for path in finalTrain:
                f.write(f'{path}\n')
        with open(self.testPath, 'w') as f:
            for path in finalTest:
                f.write(f'{path}\n')

    def getVocabFromPaths(self, paths):
        vocab = set()
        for curPath in paths:
            curPath = curPath.replace(os.sep, '/')
            clip = curPath.split('/')[-1]
            word = clip.split('_')[0]
            vocab.add(word)
        return vocab

    def getWordFromPath(self, path):
        path = path.replace(os.sep, '/')
        word = path.split('/')[-1].split('_')[0]
        return word

    def getPrefixFromPath(self, path):
        path = path.replace(os.sep, '/')
        pathLen = len(path)
        prefix = ''
        for i in range(pathLen-1,0,-1):
            if path[i] == '_':
                prefix = path[0:i]
                break
        return prefix

    def filterPaths(self, paths, vocab):
        leftPaths = []
        for curPath in paths:
            curPath = curPath.replace(os.sep, '/')
            clip = curPath.split('/')[-1]
            word = clip.split('_')[0]
            if word in vocab:
                leftPaths.append(curPath)
        return leftPaths

    def getCropMaxLen(self, markData):
        maxLen = -1
        for key in markData.keys():
            curIndex = markData[key]
            curLen = curIndex[1] - curIndex[0]
            maxLen = max(maxLen, curLen)
        return maxLen

    def changeMarkData(self, markData, maxLen):
        for key in markData.keys():
            markData[key][1] = markData[key][0] + maxLen
        return markData

    def catTestToTrain(self):
        allTrain = readDataFromFile(self.trainPath)
        allTest = readDataFromFile(self.testPath)
        res = allTrain + allTest
        with open(self.trainPath, 'w') as f:
            for path in res:
                f.write(f'{path}\n')

    def getAllPathFromAllDirList(self, dirList):
        totalPath = []
        for dirPath in dirList:
            allPath = getAllPathFromDir(dirPath, self.fileFormat)
            totalPath += allPath
        return totalPath

    def savePathToFile(self, path, data):
        with open(path, 'w') as f:
            for curPath in data:
                f.write(f'{curPath}\n')